#include "pxt.h"
#include "esp_log.h"
#include "platform.h"
#include "vm.h"

#if PXT_USB
#include "esp_vfs.h"
#include "tinyusb.h"
#include "tusb_msc_storage.h"
#include "tusb_cdc_acm.h"
#include "uf2hid.h"
#endif

#define LOG(msg, ...) DMESG("USB: " msg, ##__VA_ARGS__)
#define LOGV(msg, ...) ((void)0)
#define ERROR(msg, ...) DMESG("USB-ERROR: " msg, ##__VA_ARGS__)

#if !PXT_USB

extern "C" void usb_init() {}

#else

// 260 bytes needed for biggest JD packets (with overheads)
#define HF2_BUF_SIZE 260
#define BASE_PATH "/usb" // base path to mount the partition

static const char *TAG = "pxt_usb";
static uint8_t rx_buf[CONFIG_TINYUSB_CDC_RX_BUFSIZE + 1];

typedef struct {
    uint16_t size;
    uint8_t serial;
    union {
        uint8_t buf[HF2_BUF_SIZE];
        uint32_t buf32[HF2_BUF_SIZE / 4];
        uint16_t buf16[HF2_BUF_SIZE / 2];
        HF2_Command cmd;
        HF2_Response resp;
    };
} HF2_Buffer;

class HF2 {
    bool gotSomePacket;
    // uint32_t lastExchange;

  public:
    bool connected;
    HF2_Buffer pkt;

    int sendResponse(int size);
    void recv(uint8_t buf[64]);
    int sendResponseWithData(const void *data, int size);
    int sendEvent(uint32_t evId, const void *data, int size);
    void sendBuffer(uint8_t flag, const void *data, unsigned size, uint32_t prepend = -1);
    int handlePkt(int sz);
    int sendSerial(const void *data, int size, int isError = 0);

    HF2() {}
};

static HF2 hf2;

static void tinyusb_cdc_rx_callback(int itf, cdcacm_event_t *event)
{
    /* initialization */
    size_t rx_size = 0;
    uint8_t buf[CONFIG_TINYUSB_CDC_RX_BUFSIZE + 1];
    /* read */
    esp_err_t ret = tinyusb_cdcacm_read((tinyusb_cdcacm_itf_t)itf, rx_buf, CONFIG_TINYUSB_CDC_RX_BUFSIZE, &rx_size);
    if (ret == ESP_OK) {
        hf2.recv(buf);
        ESP_LOGI(TAG, "Received %d bytes", rx_size);
    } else {
        ESP_LOGE(TAG, "Read Error");
    }
}

static esp_err_t storage_init_spiflash(wl_handle_t *wl_handle)
{
    ESP_LOGI(TAG, "Initializing wear levelling");

    const esp_partition_t *data_partition = esp_partition_find_first(ESP_PARTITION_TYPE_DATA, ESP_PARTITION_SUBTYPE_DATA_FAT, NULL);
    if (data_partition == NULL) {
        ESP_LOGE(TAG, "Failed to find FATFS partition. Check the partition table.");
        return ESP_ERR_NOT_FOUND;
    }

    return wl_mount(data_partition, wl_handle);
}

static void on_cdc_line_state_changed(int itf, cdcacm_event_t *event) {
    hf2.connected = event->line_state_changed_data.dtr && event->line_state_changed_data.rts;
    LOG("connected: %d", hf2.connected);
}

static void tinyusb_cdc_line_state_changed_callback(int itf, cdcacm_event_t *event)
{
    int dtr = event->line_state_changed_data.dtr;
    int rts = event->line_state_changed_data.rts;
    ESP_LOGI(TAG, "Line state changed on channel %d: DTR:%d, RTS:%d", itf, dtr, rts);
}

void HF2::recv(uint8_t buf[64]) {
    uint8_t tag = buf[0];
    if (pkt.size && (tag & HF2_FLAG_SERIAL_OUT)) {
        ERROR("serial in middle of cmd");
        return;
    }

    int size = tag & HF2_SIZE_MASK;
    if (pkt.size + size > (int)sizeof(pkt.buf)) {
        ERROR("pkt too large");
        return;
    }

    memcpy(pkt.buf + pkt.size, buf + 1, size);
    pkt.size += size;
    tag &= HF2_FLAG_MASK;
    if (tag != HF2_FLAG_CMDPKT_BODY) {
        if (tag == HF2_FLAG_CMDPKT_LAST)
            pkt.serial = 0;
        else if (tag == HF2_FLAG_SERIAL_OUT)
            pkt.serial = 1;
        else
            pkt.serial = 2;
        int sz = pkt.size;
        pkt.size = 0;
        handlePkt(sz);
    }
}

const char *uf2_info() {
    return "ESP32-S2";
}

void reboot_to_uf2(void);

static void jdLog(const uint8_t *frame) {
    hf2.sendEvent(HF2_EV_JDS_PACKET, frame, frame[2] + 12);
}

int HF2::handlePkt(int sz) {
    if (pkt.serial) {
        // TODO raise some event?
        return 0;
    }

    LOGV("HF2 sz=%d CMD=%x", sz, pkt.buf32[0]);

    // one has to be careful dealing with these, as they share memory
    HF2_Command *cmd = &pkt.cmd;
    HF2_Response *resp = &pkt.resp;

    uint32_t cmdId = cmd->command_id;
    resp->tag = cmd->tag;
    resp->status16 = HF2_STATUS_OK;

    //#define checkDataSize(str, add) assert(sz == 8 + (int)sizeof(cmd->str) + (int)(add))

    // lastExchange = current_time_ms();
    gotSomePacket = true;

    switch (cmdId) {
    case HF2_CMD_INFO:
        return sendResponseWithData(uf2_info(), strlen(uf2_info()));

    case HF2_CMD_BININFO:
        resp->bininfo.mode = HF2_MODE_USERSPACE;
        resp->bininfo.flash_page_size = 0;
        resp->bininfo.flash_num_pages = 0;
        resp->bininfo.max_message_size = sizeof(pkt.buf);
        resp->bininfo.uf2_family = 0xbfdd4eee;
        return sendResponse(sizeof(resp->bininfo));

    case HF2_CMD_RESET_INTO_APP:
        target_reset();
        break;

    case HF2_CMD_RESET_INTO_BOOTLOADER:
        reboot_to_uf2();
        break;

    case HF2_CMD_DMESG:
        // TODO
        break;

    case HF2_CMD_JDS_CONFIG:
        if (cmd->data8[0]) {
            pxt::logJDFrame = jdLog;
        } else {
            pxt::logJDFrame = NULL;
        }
        return sendResponse(0);

    case HF2_CMD_JDS_SEND:
        if (pxt::sendJDFrame) {
            pxt::sendJDFrame(cmd->data8);
            return sendResponse(0);
        } else {
            resp->status16 = HF2_STATUS_INVALID_STATE;
            return sendResponse(0);
        }

    default:
        // command not understood
        resp->status16 = HF2_STATUS_INVALID_CMD;
        break;
    }

    return sendResponse(0);
}

struct BufferEntry {
    unsigned size;
    uint8_t flag;
    uint8_t data[0];
};

static void send_buffer_core(void *ent_) {
    auto ent = (BufferEntry *)ent_;

    uint32_t buf[64 / 4]; // aligned
    auto size = ent->size;
    auto data = ent->data;

    while (hf2.connected && size > 0) {
        memset(buf + 1, 0, 60);
        int s = 63;
        if (size <= 63) {
            s = size;
            buf[0] = ent->flag;
        } else {
            buf[0] = ent->flag == HF2_FLAG_CMDPKT_LAST ? HF2_FLAG_CMDPKT_BODY : ent->flag;
        }
        buf[0] |= s;
        uint8_t *dst = (uint8_t *)buf;
        dst++;
        memcpy(dst, data, s);
        data = data + s;
        size -= s;

        if (tinyusb_cdcacm_write_queue(TINYUSB_CDC_ACM_0, (uint8_t *)buf, sizeof(buf)) <
            sizeof(buf))
            DMESG("CDC write fail");
        // tinyusb_cdcacm_write_flush(TINYUSB_CDC_ACM_0, 0); - prints warnings
    }

    xfree(ent);
}

void HF2::sendBuffer(uint8_t flag, const void *data, unsigned size, uint32_t prepend) {
    if (!connected)
        return;

    if (prepend + 1)
        size += 4;

    auto ent = (BufferEntry *)xmalloc(sizeof(BufferEntry) + size);
    ent->size = size;
    ent->flag = flag;
    auto dst = ent->data;

    if (prepend + 1) {
        memcpy(dst, &prepend, 4);
        dst += 4;
        size -= 4;
    }

    memcpy(dst, data, size);

    if (worker_run_wait(fg_worker, send_buffer_core, ent))
        DMESG("HF2 queue full");
}

int HF2::sendEvent(uint32_t evId, const void *data, int size) {
    sendBuffer(HF2_FLAG_CMDPKT_LAST, data, size, evId);
    return 0;
}

int HF2::sendSerial(const void *data, int size, int isError) {
    if (!connected)
        return 0;

    sendBuffer(isError ? HF2_FLAG_SERIAL_ERR : HF2_FLAG_SERIAL_OUT, data, size);

    return 0;
}

int HF2::sendResponse(int size) {
    sendBuffer(HF2_FLAG_CMDPKT_LAST, pkt.buf, 4 + size);
    return 0;
}

int HF2::sendResponseWithData(const void *data, int size) {
    if (size <= (int)sizeof(pkt.buf) - 4) {
        memcpy(pkt.resp.data8, data, size);
        return sendResponse(size);
    } else {
        sendBuffer(HF2_FLAG_CMDPKT_LAST, data, size, pkt.resp.eventId);
        return 0;
    }
}

#endif

namespace pxt {
void sendSerial(const char *data, int len) {
    printf(LOG_BOLD(LOG_COLOR_PURPLE) "%s" LOG_RESET_COLOR, data);
#if PXT_USB
    hf2.sendSerial(data, len);
#endif
}
} // namespace pxt

// https://gist.github.com/brgaulin/2dec28baf5e9e11dfd7ef8354adf103d


extern "C" {

extern unsigned char built_binary_pxt64[];
extern unsigned int built_binary_pxt64_len;

FileEntry entries[MAX_FILES];
static int imageState = -1;
static uint8_t * usb_app_image;
// static uint8_t * usb_app_image_aligned = NULL;
static uint8_t * usb_app_image_write = NULL;
static uint32_t image_size = 0;
static uint32_t image_written = 0;
static void swap(FileEntry &a, FileEntry &b) {
    FileEntry temp = a;
    a = b;
    b = temp;
}

static void sortFileEntries(FileEntry *entries, int count) {
    for (int i = 0; i < count - 1; ++i) {
        for (int j = 0; j < count - 1 - i; ++j) {
            if (entries[j].modification_time < entries[j + 1].modification_time) {
                swap(entries[j], entries[j + 1]);
            }
        }
    }
}

uint8_t * usb_load_pxt(const char * name){
    if (name == NULL) {
        ESP_LOGI(TAG, "Loaded built-in binary: %p %d", built_binary_pxt64, built_binary_pxt64_len);
        memcpy(usb_app_image, built_binary_pxt64, built_binary_pxt64_len);
        return usb_app_image;
    }
    char fullpath[PATH_MAX];
    snprintf(fullpath, sizeof(fullpath), "%s/%s", BASE_PATH, name);
    FILE *file = fopen(fullpath, "rb");
    if (!file) {
        ESP_LOGE(TAG, "Failed to open file: %s", fullpath);
        return NULL;
    }

    fseek(file, 0, SEEK_END);
    long filesize = ftell(file);
    fseek(file, 0, SEEK_SET);

    if (filesize <= 0) {
        fclose(file);
        ESP_LOGE(TAG, "Invalid file size: %ld", filesize);
        return NULL;
    }

    // if (usb_app_image)
    //     free(usb_app_image);

    // usb_app_image = (uint8_t *)heap_caps_malloc(filesize+256, MALLOC_CAP_SPIRAM);
    // if (!usb_app_image) {
    //     fclose(file);
    //     ESP_LOGE(TAG, "Failed to allocate memory");
    //     return NULL;
    // }

    // align to 16 bytes
    // usb_app_image_aligned = (uint8_t *)(((uintptr_t)usb_app_image + 15) & ~15);

    size_t read_size = fread(usb_app_image, 1, filesize, file);
    fclose(file);

    if (read_size != filesize) {
        // free(usb_app_image);
        // usb_app_image = NULL;
        ESP_LOGE(TAG, "Failed to read file content");
        return NULL;
    }

    ESP_LOGI(TAG, "Successfully loaded file: %s, size: %d", name, filesize);

    return usb_app_image;

}

int usb_pre_mount(){

    // list all files
    DIR *dir = opendir(BASE_PATH);
    char fullpath[PATH_MAX];
    int count = 0;
    memset(entries, 0, sizeof(entries));

    if (dir) {
        struct dirent *de;
        struct stat file_stat;

        while ((de = readdir(dir)) != NULL) {
            ESP_LOGI(TAG, "File: %s", de->d_name);
            if (strstr(de->d_name, ".pxt64") != NULL && count < MAX_FILES) {
                // Construct full file path
                snprintf(fullpath, sizeof(fullpath), "%s/%s", BASE_PATH, de->d_name);
                ESP_LOGI(TAG, "Full path: %s", fullpath);

                // Get file stats
                int ret = stat(fullpath, &file_stat);
                if (ret == 0) {
                    ESP_LOGI(TAG, "File size: %ld", file_stat.st_size);
                    ESP_LOGI(TAG, "File last modified: %s", ctime(&file_stat.st_mtime));
                    entries[count].name = strdup(de->d_name);
                    entries[count].modification_time = file_stat.st_mtime;
                    count++;
                } else {
                    ESP_LOGE(TAG, "Failed to get file stats");
                }
            }
        }
        closedir(dir);
        sortFileEntries(entries, count);
        // for (int i = 0; i < count; i++) {
        //     ESP_LOGI(TAG, "File: %s, Last modified: %s", entries[i].name, ctime(&entries[i].modification_time));
        // }
        return ESP_OK;
    } else {
        return ESP_ERR_NOT_FOUND;
    }
}

void usb_write_cb(uint8_t * buffer, uint32_t bufsize, uint32_t lba){
    // ESP_LOGI(TAG, "Write CB: %p LBA: %d, Size: %d", buffer, lba, bufsize);
    // check header
    if (lba > 100 && bufsize > 0){
        VMImageSection *sect = (VMImageSection *)buffer;
        if (sect->type == pxt::SectionType::InfoHeader){
            imageState = 3;
            auto hd = (VMImageHeader *)sect->data;
            if (hd->magic0 == VM_MAGIC0 || hd->magic1 == VM_MAGIC1){
                ESP_LOGI(TAG, "Valid image header found sz=%d", hd->imageSize);
                // if (usb_app_image)
                //     free(usb_app_image);
                // usb_app_image = (uint8_t *)heap_caps_malloc(hd->imageSize + 64, MALLOC_CAP_DEFAULT | MALLOC_CAP_SPIRAM);
                // if (usb_app_image){
                    // usb_app_image_aligned = (uint8_t *)(((uintptr_t)usb_app_image + 15) & ~15);
                usb_app_image_write = usb_app_image;
                image_size = hd->imageSize * 1.1;
                image_written = 0;
                imageState = 2;
                panicCode = 952;
                // } else {
                //     ESP_LOGE(TAG, "Failed to allocate memory for app");
                //     usb_app_image_write = NULL;
                // }
            }
        }
        if (usb_app_image_write){
            memcpy(usb_app_image_write, buffer, bufsize);
            usb_app_image_write += bufsize;
            image_written += bufsize;
            if (image_written >= image_size){
                imageState = 1;
            }
        }
    } else if (lba < 100 && bufsize == 0){
        // write stop
        // issue panic and reboot vm to load new image
        panicCode = 953;
    }
}

void usb_init() {
    LOG("init");
    usb_app_image = (uint8_t *)heap_caps_malloc(1024 * 1024, MALLOC_CAP_SPIRAM | MALLOC_CAP_32BIT);
    static wl_handle_t wl_handle = WL_INVALID_HANDLE;
    ESP_ERROR_CHECK(storage_init_spiflash(&wl_handle));

    const tinyusb_msc_spiflash_config_t config_spi = {
        .wl_handle = wl_handle,
        .panic_code = &panicCode,
        .write_callback = &usb_write_cb,
    };
    ESP_ERROR_CHECK(tinyusb_msc_storage_init_spiflash(&config_spi));
    ESP_ERROR_CHECK(tinyusb_msc_storage_mount(BASE_PATH));


    LOG("USB Composite initialization");
    const tinyusb_config_t tusb_cfg = {
        .device_descriptor = NULL,
        .string_descriptor = NULL,
        .string_descriptor_count = 0,
        .external_phy = false,
#if (TUD_OPT_HIGH_SPEED)
        .fs_configuration_descriptor = NULL,
        .hs_configuration_descriptor = NULL,
        .qualifier_descriptor = NULL,
#else
        .configuration_descriptor = NULL,
#endif // TUD_OPT_HIGH_SPEED
    };
    ESP_ERROR_CHECK(tinyusb_driver_install(&tusb_cfg));

    tinyusb_config_cdcacm_t acm_cfg = {
        .usb_dev = TINYUSB_USBDEV_0,
        .cdc_port = TINYUSB_CDC_ACM_0,
        .rx_unread_buf_sz = 64,
        .callback_rx = &tinyusb_cdc_rx_callback, // the first way to register a callback
        .callback_rx_wanted_char = NULL,
        .callback_line_state_changed = NULL,
        .callback_line_coding_changed = NULL
    };

    ESP_ERROR_CHECK(tusb_cdc_acm_init(&acm_cfg));
    /* the second way to register a callback */
    ESP_ERROR_CHECK(tinyusb_cdcacm_register_callback(
                        TINYUSB_CDC_ACM_0,
                        CDC_EVENT_LINE_STATE_CHANGED,
                        &tinyusb_cdc_line_state_changed_callback));


    LOG("init done");
}

uint8_t * usb_get_new_image(){
    return usb_app_image;
}

int usb_get_write_progress(){
    int n = image_written * 100 / image_size;
    return n;
}

} // extern "C"
